import warnings
warnings.filterwarnings('ignore')

import util_for_output_zh

import pandas as pd
import numpy as np
from pyswarm import pso

import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np

from util_set_zh_matplot import plt

# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)


# 设置显示选项
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', 100)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)

'''额外的约束条件 需要 calc_pass_probability(...) > 0.8
from scipy.stats import norm

def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        -1.260742 
        + 0.076338 * week + 0.076884 * bmi 
        - 0.005221 * week**2 - 0.002285 * bmi**2 + 0.000798 * week * bmi 
        + 0.000138 * week**3 + 0.000016 * bmi**3 - 0.000063 * week**2 * bmi + 0.000022 * week * bmi**2
    )

def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        0.0524
        + 0.1736 * week
        + 0.1089 * bmi
        + -0.4638 * week**2
        + 0.0131 * week * bmi
        + -0.2792 * bmi**2
        + 0.3613 * week**3
        + -0.1273 * week**2 * bmi
        + 0.1259 * week * bmi**2
        + 0.1202 * bmi**3
    )
    
def calc_pass_probability(week, bmi, sigma):
    """计算达标概率 P(Yᵢ ≥ 0.04)"""
    y_hat = predict_y(week, bmi)
    z_score = (0.04 - y_hat) / sigma
    pass_prob = 1 - norm.cdf(z_score)
    return pass_prob
'''
from scipy.stats import norm

import os
out_dir = 'ques2_wucha'
os.makedirs(out_dir , exist_ok=True)

# def predict_y(week, bmi):
#     """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
#     return (
#         -1.260742 
#         + 0.076338 * week + 0.076884 * bmi 
#         - 0.005221 * week**2 - 0.002285 * bmi**2 + 0.000798 * week * bmi 
#         + 0.000138 * week**3 + 0.000016 * bmi**3 - 0.000063 * week**2 * bmi + 0.000022 * week * bmi**2
#     )
def predict_y(week, bmi):
    """根据拟合公式计算模型预测的浓度值Ŷᵢ"""
    return (
        0.0524
        + 0.1736 * week
        + 0.1089 * bmi
        + -0.4638 * week**2
        + 0.0131 * week * bmi
        + -0.2792 * bmi**2
        + 0.3613 * week**3
        + -0.1273 * week**2 * bmi
        + 0.1259 * week * bmi**2
        + 0.1202 * bmi**3
    )

# def calc_pass_probability(week, bmi, sigma):
#     """计算达标概率 P(Yᵢ ≥ 0.04)"""
#     y_hat = predict_y(week, bmi)
#     z_score = (0.04 - y_hat) / sigma
#     pass_prob = 1 - norm.cdf(z_score)
#     return pass_prob
def calc_pass_probability(week, bmi, sigma):
    """计算达标概率 P(Yᵢ ≥ 0.04)"""
    y_hat = predict_y(week, bmi)
    # 添加合理性检查
    if sigma < 0.01:  # 防止除零错误
        sigma = 0.01
    z_score = (0.04 - y_hat) / sigma
    pass_prob = 1 - norm.cdf(z_score)
    return max(0.0, min(1.0, pass_prob))  # 限制在[0,1]范围内

# 1. 数据加载与预处理
def convert_gestational_week(week_str):
    """快速转换孕周为数值"""
    if 'w+' in week_str:
        w, d = week_str.split('w+')
        return float(w) + float(d)/7
    return float(week_str.split('w')[0])

# 加载数据
df = pd.read_csv('ques1_prepare_data/sheet1_男胎检测数据_fillnan_sifted_outlier.csv')
df = df[['孕妇BMI', '检测孕周', 'Y染色体浓度']]
df = df[df['Y染色体浓度'] > 0.04].copy()
df['检测孕周数值'] = df['检测孕周'].apply(convert_gestational_week)
df = df.drop('检测孕周', axis=1)


# 2. 动态BMI分段
# 动态BMI分段改进版
def dynamic_bmi_segmentation(data, n_segments=4, min_samples=50):
    bmi = data['孕妇BMI'].sort_values().values
    n = len(bmi)
    
    # 确保每个分段有足够样本
    cutoffs = []
    for i in range(1, n_segments):
        pos = min_samples * i
        if pos < n:
            cutoffs.append(bmi[pos])
    
    # 合并过于接近的分界点
    cutoffs = np.unique(cutoffs)
    return cutoffs
# def dynamic_bmi_segmentation(data, n_segments=4, min_ratio=0.1):
#     bmi = data['孕妇BMI'].values
#     total = len(data)
#     min_samples = total * min_ratio
    
#     quantiles = np.linspace(0, 1, n_segments + 1)[1:-1]
#     cutoffs = np.percentile(bmi, quantiles * 100)
    
#     segments = np.digitize(bmi, cutoffs)
#     counts = np.bincount(segments, minlength=n_segments)
    
#     if not all(counts >= min_samples):
#         sorted_bmi = np.sort(bmi)
#         cutoffs = []
#         for i in range(1, n_segments):
#             cutoff_idx = int(min_samples * i)
#             cutoffs.append(sorted_bmi[cutoff_idx])
#         cutoffs = np.unique(cutoffs)
    
#     return np.sort(cutoffs)

# 重新分段
# bmi_cutoffs = dynamic_bmi_segmentation(df, n_segments=4, min_samples=50)
# print("调整后的BMI分界点:", bmi_cutoffs)
def balanced_bmi_segmentation(data, min_samples=100, n_segments=4):
    """确保每个分段有足够且均衡的样本量"""
    bmi = data['孕妇BMI'].sort_values().values
    n = len(bmi)
    
    # 计算初始分位数
    quantiles = np.linspace(0, 1, n_segments + 1)[1:-1]
    cutoffs = np.percentile(bmi, quantiles * 100)
    
    # 调整分界点确保最小样本量
    adjusted_cutoffs = []
    pos = 0
    for i in range(n_segments - 1):
        new_pos = min(pos + min_samples, n - min_samples * (n_segments - i - 1))
        adjusted_cutoffs.append(bmi[new_pos])
        pos = new_pos
    
    return np.unique(adjusted_cutoffs)
n_segments = 4
bmi_cutoffs = balanced_bmi_segmentation(df, min_samples=200, n_segments=4)
print("调整后的BMI分界点:", bmi_cutoffs)

# 预计算每个BMI分段的数据
seg_data_list = []
bmi_ranges = []
for i in range(n_segments):
    if i == 0:
        seg_data = df[df['孕妇BMI'] < bmi_cutoffs[0]]
        bmi_range = (df['孕妇BMI'].min(), bmi_cutoffs[0])
    elif i == n_segments - 1:
        seg_data = df[df['孕妇BMI'] >= bmi_cutoffs[-1]]
        bmi_range = (bmi_cutoffs[-1], df['孕妇BMI'].max())
    else:
        seg_data = df[(df['孕妇BMI'] >= bmi_cutoffs[i-1]) & 
                     (df['孕妇BMI'] < bmi_cutoffs[i])]
        bmi_range = (bmi_cutoffs[i-1], bmi_cutoffs[i])
    seg_data_list.append(seg_data)
    bmi_ranges.append(bmi_range)

# 在数据预处理阶段，为每个BMI分段计算标准差
seg_std_list = []
for seg_data in seg_data_list:
    if len(seg_data) > 1:
        seg_std_list.append(seg_data['Y染色体浓度'].std())
    else:
        seg_std_list.append(0.05)  # 默认值，当分段样本不足时使用


# 3. 定义时间权重函数
def get_time_weight(gestational_week):
    """根据孕周返回权重分数"""
    if gestational_week < 11: return 2.0
    elif 11 <= gestational_week <= 12: return 1.75
    elif 12 < gestational_week <= 15: return 1.5
    elif 15 < gestational_week <= 17: return 1.4
    elif 17 < gestational_week <= 19: return 1.3
    elif 19 < gestational_week <= 21: return 1.2
    elif 21 < gestational_week <= 23: return 1.1
    elif 23 < gestational_week <= 25: return 1.0
    elif 25 < gestational_week <= 27: return 0.5
    else: return 0.0

# 4. 优化的目标函数
def constrained_objective(x):
    """基于时间权重的新目标函数"""
    # 约束检查
    if not (x[0] <= x[1] <= x[2] <= x[3]):
        return -1e6  # 不可行解
    
    if any(x < 10) or any(x > 30):
        return -1e6
    
    # 检查达标概率约束
    for i in range(n_segments):
        bmi_range = bmi_ranges[i]
        avg_bmi = np.mean(bmi_range)  # 使用分段BMI的平均值计算
        # 使用该分段的标准差
        if calc_pass_probability(x[i], avg_bmi, seg_std_list[i]) <= 0.8:
            return -1e6  # 不满足达标概率约束
    
    total_score = 0.0
    diversity_penalty = 0.0
    
    # 计算每个分段的得分
    for i in range(n_segments):
        seg_data = seg_data_list[i]
        if len(seg_data) == 0:
            continue
        
        # 计算该分段在建议时间前的样本
        early_samples = seg_data[seg_data['检测孕周数值'] <= x[i]]
        
        # 计算这些样本的权重得分总和
        segment_score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
        
        # 标准化得分（除以总样本量）
        normalized_score = segment_score / len(seg_data)
        
        total_score += normalized_score
    
    # 多样性惩罚（防止分段过于接近）
    for i in range(1, n_segments):
        diff = x[i] - x[i-1]
        if diff < 1.0:  # 最小允许差异1周
            diversity_penalty += (1.0 - diff) * 100
    
    # 我们希望最大化总得分，所以返回负值（因为PSO默认最小化）
    return -(total_score - diversity_penalty)

# 5. 运行PSO优化
lb = [10, 12, 14, 16]  # 下限
ub = [20, 22, 24, 25]  # 上限

best_x, best_f = pso(
    constrained_objective,
    lb, ub,
    swarmsize=30,
    maxiter=150,
    debug=True,
    phip=0.7,
    phig=0.8,
    minstep=0.1,
    minfunc=0.1
)

# 6. 结果分析
print("\n=== 优化结果 ===")
for i in range(n_segments):
    lower = bmi_cutoffs[i-1] if i > 0 else -np.inf
    upper = bmi_cutoffs[i] if i < n_segments - 1 else np.inf
    seg_data = seg_data_list[i]
    
    print(f"BMI分段 {lower:.1f}-{upper:.1f}:")
    print(f"  最佳检测时间: {best_x[i]:.2f} 周")
    
    # 计算该分段的得分
    early_samples = seg_data[seg_data['检测孕周数值'] <= best_x[i]]
    score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
    normalized_score = score / len(seg_data)
    
    print(f"  得分: {normalized_score:.4f}")
    print(f"  样本量: {len(seg_data)}")
    print(f"  早期样本量: {len(early_samples)}")
    print(f"  早期比例: {len(early_samples)/len(seg_data):.2%}")

# 7. 可视化
plt.figure(figsize=(14, 7))
colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']

for i in range(n_segments):
    seg_data = seg_data_list[i]
    plt.scatter(
        seg_data['检测孕周数值'], 
        seg_data['孕妇BMI'],
        color=colors[i],
        alpha=0.6,
        label=f'BMI分段{i+1}'
    )
    
    plt.axvline(
        x=best_x[i],
        color=colors[i],
        linestyle='--',
        linewidth=2,
        label=f'分段{i+1}最佳时间({best_x[i]:.1f}周)'
    )

plt.xlabel('检测孕周', fontsize=12)
plt.ylabel('孕妇BMI', fontsize=12)
plt.title('基于时间权重评分的最佳检测时间点优化', fontsize=14)
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join(out_dir , 'ques2_optimized_detection_time_with_score.png'), dpi=300, bbox_inches='tight')
plt.close()
# ==================== 改进的检测误差分析部分 ====================
def improved_error_analysis():
    """改进的误差分析方法"""
    # 1. 临床参数初始化
    device_profiles = {
        'Device_A': {'sigma': 0.15, 'precision': 0.98},
        'Device_B': {'sigma': 0.3, 'precision': 0.95}
    }
    sigma_levels = np.linspace(0.1, 0.5, 5)
    sigma_range = np.linspace(0.1, 0.5, 5)  # 典型NIPT检测误差0.1-0.5周
    sigma_multipliers = [0.5, 0.8, 1.0, 1.2, 1.5]  # 误差波动范围
    base_std = np.mean(seg_std_list)  # 使用平均标准差作为基准
    
    def calc_pass_probability(week, avg_bmi, sigma, n_samples=100000):
        """
        基于真实临床数据的概率计算：
        改进点：
        1. 增加BMI对成功率的负面影响
        2. 引入误差累积效应
        3. 调整孕周成功率曲线，使其更符合实际
        """
        # 基础误差模型
        true_weeks = np.random.normal(week, sigma, n_samples)
        
        # 1. 孕周越大检测成功率越高（调整曲线使其更合理）
        # 新的S型曲线：10周约40%，15周约70%，20周约90%，25周约95%
        base_success_rate = 0.95 / (1 + np.exp(-0.3 * (week - 15)))
        
        # 2. BMI对成功率的负面影响（BMI越高，成功率越低）
        bmi_penalty = 1.0 - (max(0, avg_bmi - 20) * 0.02)
        bmi_penalty = np.clip(bmi_penalty, 0.6, 1.0)  # 最低保留60%成功率
        
        # 3. 误差导致的失败概率（误差越大，失败概率越高）
        error_magnitude = np.abs(true_weeks - week) / week  # 相对误差
        error_penalty = np.mean(0.5 * sigma + 0.3 * error_magnitude)
        
        # 综合达标概率（多重因素影响）
        combined_prob = base_success_rate * bmi_penalty * (1 - error_penalty)
        return np.clip(combined_prob, 0.3, 0.98)  # 限制概率范围，更符合实际情况

    sensitivity_results = []
    robustness_results = []
    
    # 3. 敏感性分析：不同标准差对达标概率的影响
    plt.figure(figsize=(12, 6))
    for i, bmi_range in enumerate(bmi_ranges):
        avg_bmi = np.mean(bmi_range)
        pass_probs = []
        
        for sigma in sigma_range:
            # 计算在当前标准差下的达标概率
            pass_prob = calc_pass_probability(best_x[i], avg_bmi, sigma)
            pass_probs.append(pass_prob)
            
            # 记录结果
            sensitivity_results.append({
                'bmi_segment': i,
                'bmi_range': f"{bmi_range[0]:.1f}-{bmi_range[1]:.1f}",
                'sigma': sigma,
                'pass_probability': pass_prob,
                'optimal_week': best_x[i]
            })
        
        # 绘制曲线
        plt.plot(sigma_range, pass_probs, 
                 label=f'BMI分段 {bmi_range[0]:.1f}-{bmi_range[1]:.1f}',
                 marker='o')

    plt.axhline(y=0.8, color='r', linestyle='--', label='达标阈值 (0.8)')
    plt.xlabel('标准差(σ)', fontsize=12)
    plt.ylabel('达标概率', fontsize=12)
    plt.title('不同标准差下的达标概率变化', fontsize=14)
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'improved_sensitivity_analysis.png'), dpi=300, bbox_inches='tight')
    plt.close()

    # 4. 稳健性分析（关键修改）
    for sigma_mult in sigma_multipliers:
        modified_std = base_std * sigma_mult
        
        for i in range(n_segments):
            avg_bmi = np.mean(bmi_ranges[i])
            seg_data = seg_data_list[i]
            
            def objective(x):
                week = x[0]
                if not (10 <= week <= 25):
                    return 1e6
                
                prob = calc_pass_probability(week, avg_bmi, modified_std)
                if prob < 0.8:
                    return 1e6 + (0.8 - prob)*1000
                
                early_samples = seg_data[seg_data['检测孕周数值'] <= week]
                if len(early_samples) == 0:
                    return 1e6
                score = early_samples['检测孕周数值'].apply(get_time_weight).mean()
                return -score
            
            def segment_objective(x):
                week = x[0]
                # 硬性约束条件
                if not (10 <= week <= 20):  # 限制在更合理的10-20周
                    return 1e8  # 大幅提高惩罚
                
                prob = calc_pass_probability(week, avg_bmi, modified_std)
                if prob < 0.95:  # 提高达标标准至95%
                    return 1e6 + (0.95 - prob)*1e4
                
                # 计算得分时考虑临床优先级
                early_samples = seg_data[seg_data['检测孕周数值'] <= week]
                score = np.sum([
                    get_time_weight(w) * clinical_priority(avg_bmi) 
                    for w in early_samples['检测孕周数值']
                ])
                return -score  # 最小化负得分

            def clinical_priority(bmi):
                """BMI越高临床优先级越高"""
                return 1.0 + 0.05 * max(0, bmi - 25)  # BMI>25时每增加1个单位权重+5%
            
            x_opt, _ = pso(
                segment_objective,
                lb=[10], ub=[20],  # 收紧搜索范围
                swarmsize=50,      # 增加粒子数量
                maxiter=100,       # 增加迭代次数
                omega=0.7,         # 动态惯性权重
                phip=1.2,          # 增强个体经验引导
                phig=1.2,          # 增强群体经验引导
                minfunc=1e-4,      # 更高精度要求
                debug=True
            )
            
            # 记录结果
            robustness_results.append({
                'sigma_multiplier': sigma_mult,
                'actual_sigma': modified_std,
                'bmi_segment': i,
                'bmi_range': f"{bmi_ranges[i][0]:.1f}-{bmi_ranges[i][1]:.1f}",
                'original_week': best_x[i],
                'adjusted_week': x_opt[0],
                'time_change': x_opt[0] - best_x[i],
                'pass_probability': calc_pass_probability(x_opt[0], avg_bmi, modified_std),
                'original_prob': calc_pass_probability(best_x[i], avg_bmi, modified_std)
            })
            
    # 5. 保存结果
    sensitivity_df = pd.DataFrame(sensitivity_results)
    robustness_df = pd.DataFrame(robustness_results)
    
    sensitivity_df.to_csv(os.path.join(out_dir, 'improved_sensitivity_results.csv'), index=False)
    robustness_df.to_csv(os.path.join(out_dir, 'improved_robustness_results.csv'), index=False)

    # 6. 可视化稳健性分析结果
    plt.figure(figsize=(14, 8))
    
    # 时间变化图
    plt.subplot(2, 1, 1)
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        plt.plot(seg_data['sigma_multiplier'], seg_data['time_change'],
                 marker='o', label=f'分段 {i+1}')
    plt.axhline(y=0, color='k', linestyle='--')
    plt.xlabel('标准差倍数')
    plt.ylabel('时间变化(周)')
    plt.title('误差对检测时间的影响')
    plt.legend()
    plt.grid(True, alpha=0.3)

    # 达标概率图
    plt.subplot(2, 1, 2)
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        plt.plot(seg_data['sigma_multiplier'], seg_data['pass_probability'],
                 marker='o', label=f'分段 {i+1}')
    plt.axhline(y=0.8, color='r', linestyle='--', label='达标阈值')
    plt.xlabel('标准差倍数')
    plt.ylabel('达标概率')
    plt.title('误差对达标概率的影响')
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'improved_robustness_analysis.png'), dpi=300, bbox_inches='tight')
    plt.close()

    # 7. 输出关键发现
    print("\n=== 改进的误差分析关键发现 ===")
    for i in range(n_segments):
        seg_data = robustness_df[robustness_df['bmi_segment'] == i]
        max_change = seg_data['time_change'].abs().max()
        min_prob = seg_data['pass_probability'].min()
        
        # 评估敏感度
        if max_change > 2.0:
            sensitivity = '高'
        elif max_change > 1.0:
            sensitivity = '中'
        else:
            sensitivity = '低'
        
        print(f"BMI分段 {bmi_ranges[i][0]:.1f}-{bmi_ranges[i][1]:.1f}:")
        print(f"  最大时间变化: {max_change:.2f} 周")
        print(f"  最低达标概率: {min_prob:.2%}")
        print(f"  对误差敏感度: {sensitivity}")
        
        # 检查是否存在不达标情况
        if min_prob < 0.8:
            print(f"  警告: 在σ={seg_data.loc[seg_data['pass_probability'].idxmin(), 'sigma_multiplier']:.1f}时达标概率低于80%")

# 运行改进的分析
print("\n=== 运行改进的误差分析 ===")
improved_error_analysis()

# ==================== 新增的检测误差分析部分结束 ====================

# 8. 保存结果
results = []
for i in range(n_segments):
    lower = bmi_cutoffs[i-1] if i > 0 else df['孕妇BMI'].min()
    upper = bmi_cutoffs[i] if i < n_segments - 1 else df['孕妇BMI'].max()
    seg_data = seg_data_list[i]
    early_samples = seg_data[seg_data['检测孕周数值'] <= best_x[i]]
    score = early_samples['检测孕周数值'].apply(get_time_weight).sum()
    normalized_score = score / len(seg_data)
    
    results.append({
        'BMI范围': f"{lower:.1f}-{upper:.1f}",
        '最佳检测时间(周)': best_x[i],
        '得分': normalized_score,
        '样本量': len(seg_data),
        '早期样本量': len(early_samples),
        '早期比例': len(early_samples)/len(seg_data)
    })

results_df = pd.DataFrame(results)
results_df.to_csv( os.path.join( out_dir , 'ques2_optimized_detection_results.csv') , index=False)
print("\n优化结果已保存到 ques2_optimized_detection_results.csv")